from langserve import add_routes
from fastapi import FastAPI, Request, HTTPException
from agent import ai_client
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, FileResponse

from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware

import os
import re
import uuid
import config
from tools.ai_tools import create_pdf, write_to_file

# FastAPI接口定义
app = FastAPI(
    title="Java助手 - LangChain Server",
    version="1.0",
    description="支持多轮对话和检索增强的Java问答服务",
)

# 使用slowapi进行基于客户端ip的限流配置来创建一个限流器对象
limiter = Limiter(
    key_func=get_remote_address,  # 使用客户端 IP 作为限流 key
    default_limits=["5/minute"],  # 默认限流：每分钟最多 5 次
    headers_enabled=True,  # 返回 X-RateLimit-* 头
    storage_uri="redis://localhost:6379"  # 使用内存存储（生产建议用 redis://localhost:6379）
)
app.state.limiter = limiter
# FastAPI添加异常处理器
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)  # 使用 slowapi 内置 handler
# FastAPI添加中间件（限流）
app.add_middleware(SlowAPIMiddleware)

# 跨域配置
app.add_middleware(
    # 内置跨域中间件
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# 健康检查接口
@app.get("/healthz")
async def healthz():
    return {"status": "ok"}


# 请求模型定义
class ChatRequest(BaseModel):
    session_id: str
    question: str

# 导出pdf请求模型定义
class ExportPdfRequest(BaseModel):
    title: str
    content: str
    filename: str = None

# 导出文件请求模型定义
class ExportFileRequest(BaseModel):
    filename: str
    content: str


# 安全检查
safe_name_pattern = re.compile(r'^[\w\-\.]+$')
os.makedirs(config.EXPORT_DIR, exist_ok=True)


# 聊天接口（标准响应）
@app.post("/chat")
# 限流每分钟10次
@limiter.limit("10/minute")
# 异步支持，（返回的是一个python协程对象）
async def chat_endpoint(request: Request, body: ChatRequest):
    try:
        invoke_config = {"configurable": {"session_id": body.session_id}}
        result = await ai_client.CHAIN_WITH_HISTORY.ainvoke(
            {"input": body.question},
            config=invoke_config
        )
        # 如果ai返回是一个字典的话
        if isinstance(result, dict):
            answer = result.get("output") or result.get("response") or str(result)
            # 不是字典的话直接转化成字符串
        else:
            answer = str(result)
        return {"response": answer}
    except Exception as e:
        return {"error": str(e)}


# 流式聊天接口
@app.post("/chat/stream")
# 每分钟15次
@limiter.limit("15/minute")
async def chat_stream(request: Request, body: ChatRequest):
    invoke_config = {"configurable": {"session_id": body.session_id}}

    async def event_generator():
        try:
            async for chunk in ai_client.CHAIN_WITH_HISTORY.astream(
                {"input": body.question},
                config=invoke_config,
                stream_mode="values"
            ):
                # 根据 chain 返回结构处理 chunk
                text = chunk if isinstance(chunk, str) else (
                    chunk.get("output") if isinstance(chunk, dict) else str(chunk)
                )
                if text.strip():
                    yield f"data: {text}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            yield f"data: [ERROR] {str(e)}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")


# 导出 PDF
@app.post("/export/pdf")
@limiter.limit("10/minute")
async def export_pdf(request: Request, body: ExportPdfRequest):
    filename = body.filename or f"export-{uuid.uuid4().hex}.pdf"

    if not safe_name_pattern.match(filename):
        raise HTTPException(status_code=400, detail="非法文件名")

    path = os.path.abspath(os.path.join(config.EXPORT_DIR, filename))
    export_dir_abs = os.path.abspath(config.EXPORT_DIR)

    if not path.startswith(export_dir_abs):
        raise HTTPException(status_code=400, detail="非法路径")

    result = create_pdf.invoke({
        "content": body.content,
        "title": body.title,
        "filename": path
    })

    if not result or not result.get("success"):
        error_msg = result.get("error", "生成失败") if result else "生成失败"
        raise HTTPException(status_code=500, detail=error_msg)

    return {"success": True, "filename": os.path.basename(path)}


# 导出普通文件
@app.post("/export/file")
@limiter.limit("20/minute")
async def export_file(request: Request, body: ExportFileRequest):
    if not safe_name_pattern.match(body.filename):
        raise HTTPException(status_code=400, detail="非法文件名")

    path = os.path.abspath(os.path.join(config.EXPORT_DIR, body.filename))
    export_dir_abs = os.path.abspath(config.EXPORT_DIR)

    if not path.startswith(export_dir_abs):
        raise HTTPException(status_code=400, detail="非法路径")

    result = write_to_file.invoke({
        "content": body.content,
        "filename": path
    })

    if not result or not result.get("success"):
        error_msg = result.get("error", "写入失败") if result else "写入失败"
        raise HTTPException(status_code=500, detail=error_msg)

    return {"success": True, "filename": os.path.basename(path)}


# 下载文件
@app.get("/download/{filename}")
@limiter.limit("60/minute")
async def download_file(request: Request, filename: str):
    if not safe_name_pattern.match(filename):
        raise HTTPException(status_code=400, detail="非法文件名")

    path = os.path.abspath(os.path.join(config.EXPORT_DIR, filename))
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="文件不存在")

    return FileResponse(path, filename=filename)


# LangServe 路由
try:
    from langserve import add_routes

    add_routes(app, ai_client.CHAIN_WITH_HISTORY, path="/chain")
except ImportError:
    print("langserve not installed, skipping /chain route.")